import json
from elasticsearch7 import Elasticsearch, exceptions
from datetime import datetime
import pandas as pd

# 多集群配置 (支持密码认证)
CLUSTERS = [
    {
        "name": "万达集群",
        "url": "http://172.16.100.76:9200",
        "auth": ("elastic", "Elastic#9200")
    },
    {
        "name": "万虹集群",
        "url": "http://172.16.100.77:9200",
        "auth": ("elastic", "Elastic#9200")
    }
]


def get_es_client(cluster_cfg):
    """创建ES客户端连接"""
    try:
        return Elasticsearch(
            hosts=[cluster_cfg["url"]],
            http_auth=cluster_cfg.get("auth"),
            verify_certs=False,
            timeout=30,
            max_retries=2
        )
    except Exception as e:
        print(f"连接集群 {cluster_cfg['name']} 失败: {str(e)}")
        return None


def cluster_info(es):
    try:
        # 获取用户列表（需启用安全功能）
        users = es.security.get_user()
        # 获取所有索引状态
        indices = es.cat.indices(format="json", h="index,health")
        non_green = [idx["index"] for idx in indices if idx["health"] != "green"]
        # 检查集群信息
        health = es.cluster.health()

        # 获取节点资源统计
        nodes_dict = {}
        stats = es.nodes.stats(metric="os,jvm,fs")
        nodes = stats["nodes"].values()
        for node in nodes:
            node_name = {
                node["name"]: {
                    "cpu_usage": node["os"]["cpu"]["percent"],
                    "mem_usage": node["jvm"]["mem"]["heap_used_percent"],
                    "fs_total": node["fs"]["total"]["total_in_bytes"],
                    "fs_free": node["fs"]["total"]["free_in_bytes"]
                }
            }
            nodes_dict.update(node_name)

        return {
            "cluster_name": health["cluster_name"],
            "status": health["status"],
            "node_count": health["number_of_nodes"],
            "data_nodes": health["number_of_data_nodes"],
            "active_primary_shards": health["active_primary_shards"],
            "active_shards": health["active_shards"],
            "unassigned_shards": health["unassigned_shards"],
            "version": es.info()["version"]["number"],
            "user": list(users.keys()),
            "non_green": non_green,
            "node_stats": nodes_dict
        }
    except exceptions.ElasticsearchException as e:
        return {"error": str(e)}


def generate_report(report_data):
    """生成巡检报告"""
    # 创建DataFrame
    df = pd.DataFrame(report_data)

    # 列顺序调整
    columns = [
        "cluster_name", "version", "status", "node_count", "data_nodes", "active_primary_shards", "active_shards",
        "unassigned_shards", "user", "non_green", "node_stats"
    ]
    df = df[columns]

    # 生成Excel文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M")
    filename = f"ES_巡检报告_{timestamp}.xlsx"

    # 写入Excel
    with pd.ExcelWriter(filename, engine="xlsxwriter") as writer:
        df.to_excel(writer, sheet_name="集群概览", index=False)

        # 添加格式
        workbook = writer.book
        worksheet = writer.sheets["集群概览"]

        # 设置列宽
        for idx, col in enumerate(df.columns):
            width = max(len(str(col)) + 2, 15)
            worksheet.set_column(idx, idx, width)

        # 状态列颜色标记
        format_red = workbook.add_format({"bg_color": "#FFC7CE"})
        format_green = workbook.add_format({"bg_color": "#C6EFCE"})

        for row_num in range(1, len(df) + 1):
            cell_value = df.iloc[row_num - 1]["status"]
            fmt = format_green if cell_value == "green" else format_red
            worksheet.conditional_format(
                row_num, 3, row_num, 3,
                {"type": "no_errors", "format": fmt}
            )

    print(f"报告已生成: {filename}")


if __name__ == "__main__":

    report_data = []

    for cluster in CLUSTERS:
        print(f"正在检查集群: {cluster['name']}")
        es = get_es_client(cluster)
        if not es:
            print(None)
        cluster_data = cluster_info(es)
        report_data.append(cluster_data)
        # a = json.dumps(cluster_data)
        # print(a)
    generate_report(report_data)
